﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Forms;
using System.Diagnostics;
using System.Collections.Concurrent;

namespace AffinityHelper {
	public partial class Form1 : Form {
		enum AffinityType {
			PCore,
			ECore,
			All,
			Mixed
		}

		public Form1() {
			InitializeComponent();
		}

		private void Form1_Load(object sender, EventArgs e) {
			comboBox1.Text = "8";
			comboBox2.Text = "2";
			comboBox3.Text = "12";
			RefreshProcessesWithoutUsages();
			dataGridView1.Columns[0].Width = 50;
			dataGridView1.Columns[1].Width = 140;
			dataGridView1.Columns[2].Width = 160;
			dataGridView1.Columns[3].Width = 100;
			dataGridView1.RowHeadersWidth = 20;
		}

		private void button1_Click(object sender, EventArgs e) {
			try {
				DataTable table = new DataTable();
				table.Columns.Add("ID", typeof(int));
				table.Columns.Add("Name", typeof(string));
				table.Columns.Add("Age", typeof(int));
				table.Rows.Add(1, "John", 25);
				table.Rows.Add(2, "Jane", 30);
				table.Rows.Add(3, "Bob", 35);
				dataGridView1.DataSource = table;
				dataGridView1.Columns["ID"].Width = 200;
				dataGridView1.ReadOnly = true;
				dataGridView1.AllowUserToAddRows = false;
			} catch { }
		}

		private void SetAffinity(AffinityType type, int PCnt = int.MaxValue, int ECnt = int.MaxValue) {
			var selectedCells = dataGridView1.SelectedCells;
			foreach (DataGridViewCell cell in selectedCells) {
				int rowIndex = cell.RowIndex;
				int pid = int.Parse(dataGridView1.Rows[rowIndex].Cells[0].Value.ToString());
				string processName = dataGridView1.Rows[rowIndex].Cells[1].Value.ToString();
				try {
					Process process = Process.GetProcessById(pid);
					if (process.HasExited) {
						MessageBox.Show($"Process {pid} has exited");
						return;
					}

					switch (type) {
						case AffinityType.ECore:
							SetECoreAffinity(process);
							break;
						case AffinityType.PCore:
							SetPCoreAffinity(process);
							break;
						case AffinityType.All:
							SetAllAffinity(process);
							break;
						case AffinityType.Mixed:
							SetMixedAffinity(process, PCnt, ECnt);
							break;
					}
				} catch (ArgumentException ex) {
					MessageBox.Show(processName + ": " + ex.Message, "Warning", MessageBoxButtons.OK, MessageBoxIcon.Warning);
				} catch (Win32Exception ex) {
					MessageBox.Show(processName + ": " + ex.Message, "Error", MessageBoxButtons.OK, MessageBoxIcon.Error);
				}
			}
		}

		private void SetECoreAffinity(Process proc) {
			int cpuCnt = Environment.ProcessorCount;
			long affinityMask = (1L << cpuCnt) - 1;
			int doNotUseFirstCores = this.TotalPCoreLogical;
			if (doNotUseFirstCores >= cpuCnt) doNotUseFirstCores = cpuCnt / 2;
			long secondMask = affinityMask << doNotUseFirstCores;
			affinityMask &= secondMask;
			proc.ProcessorAffinity = (IntPtr)affinityMask;
		}

		private void SetPCoreAffinity(Process proc) {
			int cpuCnt = Environment.ProcessorCount;
			long affintyMask = (1L << cpuCnt) - 1;
			int useFirstCores = TotalPCoreLogical;
			long secondMask = (1L << useFirstCores) - 1;
			affintyMask &= secondMask;
			proc.ProcessorAffinity = (IntPtr)affintyMask;
		}

		private void SetAllAffinity(Process proc) {
			int cpuCnt = Environment.ProcessorCount;
			long affinityMask = (1L << cpuCnt) - 1;
			proc.ProcessorAffinity = (IntPtr)affinityMask;
		}

		private void SetMixedAffinity(Process proc, int PCnt, int ECnt) {
			int cpuCnt = Environment.ProcessorCount;
			int totalPCoreLogical = this.TotalPCoreLogical;
			int totalPCorePhysical = this.TotalPCorePhysical;
			int totalECore = cpuCnt - totalPCoreLogical;
			if (ECnt > totalECore) ECnt = totalECore;
			if (PCnt > totalPCoreLogical) PCnt = totalPCoreLogical;
			Random rd = new Random();
			HashSet<int> set = new HashSet<int>();
			long affinityMask = 0;
			while (true) {
				if (set.Count >= PCnt) break;
				int c = rd.Next(totalPCorePhysical);
				if (!set.Contains(c)) {
					set.Add(c);
					long mask;
					if (PCoreHyperThread)
						mask = 3L << (c * 2);
					else
						mask = 1L << c;
					affinityMask |= mask;
				}
			}
			if (ECnt == totalECore) {
				long totalMask = (1L << cpuCnt) - 1;
				long pMask = (1L << totalPCoreLogical) - 1;
				long eMask = totalMask ^ pMask;
				affinityMask |= eMask;
			} else {
				set.Clear();
				while (true) {
					if (set.Count >= ECnt) break;
					int c = rd.Next(totalECore);
					if (!set.Contains(c)) {
						set.Add(c);
						affinityMask |= 1L << (c + totalPCoreLogical);
					}
				}
			}
			proc.ProcessorAffinity = (IntPtr)affinityMask;
		}

		private void btnRefresh_Click(object sender, EventArgs e) {
			RefreshProcessesWithUsages();
		}

		private void RefreshProcessesWithUsages() {
			var cpuUsageList = ProcessCPUHelper.Instance.GetCPUUsages();
			dataGridView1.AutoGenerateColumns = true;

			DataTable dataTable = new DataTable();
			dataTable.Columns.Add("PID");
			dataTable.Columns.Add("Instance Name");
			dataTable.Columns.Add("% CPU Usage");
			foreach (var cpuUsage in cpuUsageList) {
				dataTable.Rows.Add(cpuUsage.pid, cpuUsage.instanceName, cpuUsage.cpuUsage);
			}
			dataGridView1.DataSource = dataTable;
		}

		private void btnRefresh2_Click(object sender, EventArgs e) {
			RefreshProcessesWithoutUsages();
			if (!string.IsNullOrEmpty(textBox1.Text)) {
				textBox1_TextChanged(sender, null);
			}
		}

		private void RefreshProcessesWithoutUsages() {
			Process[] processes = Process.GetProcesses();
			Array.Sort(processes, (x, y) => {
				if (x.ProcessName.ToLower() == "idle") return 1;
				else if (y.ProcessName.ToLower() == "idle") return -1;

				try {
					return -x.StartTime.CompareTo(y.StartTime);
				} catch (InvalidOperationException) {
					return 0;
				} catch (Win32Exception) {
					return 0;
				}
			});
			DataTable dataTable = new DataTable();
			dataTable.Columns.Add("PID");
			dataTable.Columns.Add("Instance Name");
			dataTable.Columns.Add("Start Time");
			dataTable.Columns.Add("% CPU Usage");
			foreach (Process process in processes) {
				try {
					dataTable.Rows.Add(process.Id, process.ProcessName, process.StartTime, null);
				} catch (InvalidOperationException) {
				} catch (Win32Exception) {
				}
			}
			dataGridView1.DataSource = dataTable;
		}

		private void button2_Click(object sender, EventArgs e) {
			try {
				SetAffinity(AffinityType.ECore);
			} catch (FormatException ex) {
				MessageBox.Show(ex.Message, "Warning", MessageBoxButtons.OK, MessageBoxIcon.Warning);
			}
		}

		private void button3_Click(object sender, EventArgs e) {
			try {
				SetAffinity(AffinityType.PCore);
			} catch (FormatException ex) {
				MessageBox.Show(ex.Message, "Warning", MessageBoxButtons.OK, MessageBoxIcon.Warning);
			}
		}

		private void button4_Click(object sender, EventArgs e) {
			try {
				SetAffinity(AffinityType.All);
			} catch (FormatException ex) {
				MessageBox.Show(ex.Message, "Warning", MessageBoxButtons.OK, MessageBoxIcon.Warning);
			}
		}

		private void btnMixed_Click(object sender, EventArgs e) {
			try {
				int PCnt = int.Parse(comboBox2.Text);
				int ECnt = int.Parse(comboBox3.Text);
				SetAffinity(AffinityType.Mixed, PCnt, ECnt);
			} catch (FormatException ex) {
				MessageBox.Show(ex.Message, "Warning", MessageBoxButtons.OK, MessageBoxIcon.Warning);
			}
		}

		private void btnExit_Click(object sender, EventArgs e) {
			Application.Exit();
		}

		private void textBox1_TextChanged(object sender, EventArgs e) {
			try {
				(dataGridView1.DataSource as DataTable).DefaultView.RowFilter =
					string.Format("`Instance Name` LIKE '{0}%'", textBox1.Text);
			} catch { }
		}

		private void textBox1_KeyUp(object sender, KeyEventArgs e) {

		}

		private void textBox1_KeyDown(object sender, KeyEventArgs e) {
			if (e.KeyCode == Keys.Enter && !string.IsNullOrEmpty(textBox1.Text)) {
				if (MessageBox.Show("Set as E-core affinity?", "Confirm", MessageBoxButtons.OKCancel) == DialogResult.OK)
					SetAffinity(AffinityType.ECore);
			} else if (e.KeyCode == Keys.F5) {
				btnRefresh2_Click(sender, e);
			}
		}

		private void btnSearch_Click(object sender, EventArgs e) {
			FillUsageInSearchedGrid();
		}

		private void FillUsageInSearchedGrid() {
			btnSearch.Text = "Searching";
			btnSearch.Enabled = false;
			IDictionary<string, int> namedBuffer = new Dictionary<string, int>();
			foreach (DataGridViewRow row in dataGridView1.Rows) {
				if (row.Visible) {
					int pid = int.Parse(row.Cells[0].Value.ToString());
					string name = row.Cells[1].Value.ToString();
					// float usage = ProcessCPUHelper.getUsageByIdAndName(pid, name);
					float usage = ProcessCPUHelper.getUsageByIdAndName(pid, name, namedBuffer);
					row.Cells[3].Value = usage;
				}
			}
			btnSearch.Text = "Get usage";
			btnSearch.Enabled = true;
		}

		private void Form1_Activated(object sender, EventArgs e) {
			textBox1.Focus();
		}

		public int TotalPCoreLogical {
			get {
				int PCnt = int.Parse(comboBox1.Text);
				if (PCoreHyperThread)
					PCnt *= 2;
				return PCnt;
			}
		}

		public int TotalPCorePhysical {
			get {
				int PCnt = int.Parse(comboBox1.Text);
				return PCnt;
			}
		}

		public bool PCoreHyperThread {
			get {
				return checkBoxHyperThread.Checked;
			}
		}
	}
}